package migrator

import (
	"fmt"
	"log"
	"sync"
	"sync/atomic"

	"mig/pkg/chunk"
	"mig/pkg/data"
	"mig/pkg/types"
)

func simpleMigrate(source data.Database, target data.Database,
	chunkSize int64, maxWorkers int, unequalChunks []types.UnequalChunk) ([]*types.DigestNode, error) {
	var (
		chunkNos   []int
		chunkNoMap map[int]int
	)
	migrateUnequalOnly := len(unequalChunks) > 0
	if migrateUnequalOnly {
		chunkNos = make([]int, 0, len(unequalChunks))
		chunkNoMap = make(map[int]int, len(unequalChunks))
		for ckIdx, ck := range unequalChunks {
			chunkNos = append(chunkNos, ck.ChunkNo)
			chunkNoMap[ck.ChunkNo] = ckIdx
		}
	} else {
		totalRows, err := source.Count()
		if err != nil {
			return nil, err
		}
		log.Printf("total rows: %d", totalRows)
		totalChunks := totalRows / chunkSize
		if totalRows%chunkSize > 0 {
			totalChunks++
		}
		chunkNos = make([]int, 0, totalChunks)
		for i := 0; i < int(totalChunks); i++ {
			chunkNos = append(chunkNos, i)
		}
	}
	chunkNoIdx := int32(-1)
	wg := sync.WaitGroup{}
	digests := make([]*types.DigestNode, len(chunkNos))

	mtx := sync.Mutex{}
	for i := 0; i < maxWorkers; i++ {
		wg.Add(1)
		go func() {
			var (
				ck        chunk.Chunk
				workerErr error
			)

			for {
				currentNoIdx := atomic.AddInt32(&chunkNoIdx, 1)
				if int(currentNoIdx) >= len(chunkNos) {
					wg.Done()
					return
				}
				currentNo := int64(chunkNos[currentNoIdx])
				ck, workerErr = source.ReadChunk(currentNo, chunkSize)
				if workerErr != nil {
					log.Printf("read error: %v", workerErr)
					continue
				}

				mtx.Lock()
				if migrateUnequalOnly {
					digests[chunkNoMap[int(currentNo)]] = ck.Digest()
				} else {
					digests[currentNo] = ck.Digest()
				}
				mtx.Unlock()

				if workerErr = target.WriteChunk(ck); workerErr != nil {
					log.Printf("write error: %v", workerErr)
					continue
				}
			}
		}()
	}
	wg.Wait()
	for chunkIdx, digestNode := range digests {
		digestNode.ID = fmt.Sprintf("%s_%d", source.Identifier(), chunkNos[chunkIdx])
	}
	return digests, nil
}

type CommonMigrator struct {
	chunkSize  int64
	maxWorkers int
}

func NewCommonMigrator(chunkSize int64, maxWorkers int) *CommonMigrator {
	return &CommonMigrator{
		chunkSize:  chunkSize,
		maxWorkers: maxWorkers,
	}
}

func (d *CommonMigrator) Migrate(source data.Database, target data.Database, unequalChunks []types.UnequalChunk) (*types.IntegrityInfo, error) {
	digestNodes, err := simpleMigrate(source, target, d.chunkSize, d.maxWorkers, unequalChunks)
	if err != nil {
		return nil, err
	}
	return &types.IntegrityInfo{
		DigestNodes: digestNodes,
	}, nil
}
